import tensorflow as tf
import numpy as np
from .aux_py import get_var_wrap
from .kron import kron


def ht_conv(inp,
        inp_modes,  # 输入张量[n1,n2,n3,n4]，其中n1 为l*l,n2*n3*n4=C
        out1_modes,#输出矩阵[m1,m2,m3,m4]，m1=1,m2*m3*m4=S
        out_modes,#输出矩阵乘积
        matin_ranks,  # rk
        matout_ranks,  # rk
        strides=[1, 1, 1, 1],
        padding='SAME',
        cores_initializer=tf.contrib.layers.xavier_initializer(uniform=False),
        cores_regularizer=None,
		biases_initializer=tf.constant_initializer(0.1),
		biases_regularizer=None,
        trainable=True,
        cpu_variables=False,
        scope=None):
	with tf.variable_scope(scope):
		dimin = len(inp_modes)
		dimout = len(out1_modes)
		matin_cores = []
		matout_cores = []

		cinit = cores_initializer
		creg = cores_regularizer
		for i in range(dimin//2):
			matin_cores.append(get_var_wrap('mat_corein_%d' % (i + 1),
                                            shape=[inp_modes[i]*out1_modes[i], matin_ranks[i]],
                                            initializer=cinit,
                                            regularizer=creg,
                                            trainable=trainable,
											cpu_variable=False))


		matin_kron=kron(matin_cores[0], matin_cores[1])  # Ua与Ub的克罗内克积
		for i in range(2,2+dimout//2):
			matout_cores.append(get_var_wrap('mat_coreout_%d' % (i + 1),
                                             shape=[inp_modes[i]*out1_modes[i],matout_ranks[i-2]],
                                             initializer=cinit,
                                             regularizer=creg,
                                             trainable=trainable,
											 cpu_variable=False))


		matout_kron=kron(matout_cores[0], matout_cores[1])  # Ua与Ub的克罗内克积

		blast = get_var_wrap('mat_coreinlast',
                              shape=[matin_ranks[-1]*matout_ranks[-1],1],
                              initializer=cinit,
                              regularizer=creg,
                              trainable=trainable,
							  cpu_variable=False)


		tb12 = get_var_wrap('tb12',
                              shape=[matin_ranks[0]*matin_ranks[1],matin_ranks[2]],
                              initializer=cinit,
                              regularizer=creg,
                              trainable=trainable,
							  cpu_variable=False)
		tb34 = get_var_wrap('tb34',
                              shape=[matout_ranks[0]*matout_ranks[1],matout_ranks[2]],
                              initializer=cinit,
                              regularizer=creg,
                              trainable=trainable,
							  cpu_variable=False)
		matin_kron=tf.matmul(matin_kron,tb12)
		matout_kron=tf.matmul(matout_kron,tb34)

		matout_kron=tf.reshape(matout_kron,[inp_modes[2],out1_modes[2],inp_modes[3],out1_modes[3],matout_ranks[2]])
		matout_kron=tf.transpose(matout_kron,[0,2,1,3,4])
		matout_kron=tf.reshape(matout_kron,[inp_modes[2]*inp_modes[3]*out1_modes[2]*out1_modes[3],matout_ranks[2]])
		matout_kron=tf.transpose(matout_kron,[1,0])

		blast=tf.reshape(blast,[matin_ranks[-1],matout_ranks[-1]])
		matin_kron=tf.matmul(matin_kron,blast)

		mat_cores=tf.matmul(matin_kron,matout_kron)


		mat_cores = tf.reshape(mat_cores,
                               [inp_modes[0] * out1_modes[0]*inp_modes[1], out1_modes[1], inp_modes[2] * inp_modes[3],
                                out1_modes[2] * out1_modes[3]])
		mat_cores = tf.transpose(mat_cores, [0, 2, 1, 3])
		mat_cores = tf.reshape(mat_cores, [inp_modes[0],out1_modes[0],inp_modes[1] * inp_modes[2] * inp_modes[3],
                                           out1_modes[1] * out1_modes[2] * out1_modes[3]])


		out = tf.nn.conv2d(inp, mat_cores, strides, padding)

		biases = get_var_wrap('biases',
                                  shape=[out_modes],
                                  initializer=biases_initializer,
                                  regularizer=biases_regularizer,
                                  trainable=trainable,
                                  cpu_variable=cpu_variables)
		out = tf.nn.bias_add(out, biases)
	return out
